import cv2
import numpy as np

# 全局变量
drawing = False  # 是否正在绘制
ix, iy = -1, -1  # 鼠标起始坐标
rect = (0, 0, 0, 0)  # 矩形框坐标
mask = None  # GrabCut掩码
img = None  # 原始图像
img_copy = None  # 用于绘制的图像副本

# 鼠标回调函数
def draw_rectangle(event, x, y, flags, param):
    global ix, iy, drawing, img, img_copy, mask, rect

    if event == cv2.EVENT_LBUTTONDOWN:
        drawing = True
        ix, iy = x, y
        img_copy = img.copy()  # 每次点击时重新复制原始图像

    elif event == cv2.EVENT_MOUSEMOVE:
        if drawing:
            img_copy = img.copy()
            cv2.rectangle(img_copy, (ix, iy), (x, y), (0, 255, 0), 2)

    elif event == cv2.EVENT_LBUTTONUP:
        drawing = False
        rect = (min(ix, x), min(iy, y), abs(ix - x), abs(iy - y))
        # 初始化GrabCut掩码
        mask = np.zeros(img.shape[:2], np.uint8)
        # 矩形框内设为可能前景(GC_PR_FGD)，外部设为确定背景(GC_BGD)
        mask[rect[1]:rect[1]+rect[3], rect[0]:rect[0]+rect[2]] = cv2.GC_PR_FGD
        cv2.rectangle(img_copy, (ix, iy), (x, y), (0, 255, 0), 2)
        print("矩形框已选定，按空格键开始GrabCut分割")

# 笔刷回调函数（用于修正分割结果）
def draw_brush(event, x, y, flags, param):
    global mask, img_copy

    if event == cv2.EVENT_LBUTTONDOWN:
        # 左键：标记前景
        cv2.circle(mask, (x, y), 5, cv2.GC_FGD, -1)
        cv2.circle(img_copy, (x, y), 5, (0, 0, 255), -1)
    elif event == cv2.EVENT_RBUTTONDOWN:
        # 右键：标记背景
        cv2.circle(mask, (x, y), 5, cv2.GC_BGD, -1)
        cv2.circle(img_copy, (x, y), 5, (255, 0, 0), -1)

# GrabCut分割函数
def grabcut_segmentation():
    global mask, img, img_copy

    # 初始化背景和前景模型
    bgd_model = np.zeros((1, 65), np.float64)
    fgd_model = np.zeros((1, 65), np.float64)

    # 执行GrabCut
    cv2.grabCut(img, mask, rect, bgd_model, fgd_model, 5, cv2.GC_INIT_WITH_RECT)

    # 生成二值掩码
    mask2 = np.where((mask == cv2.GC_BGD) | (mask == cv2.GC_PR_BGD), 0, 1).astype('uint8')
    # 提取前景
    result = img * mask2[:, :, np.newaxis]

    # 显示结果
    cv2.imshow('Segmentation Result', result)
    print("GrabCut分割完成！按'c'键继续修正，或按其他键退出")

# 主函数
def main():
    global img, img_copy, mask

    # 读取图像
    img = cv2.imread('input.jpg')  # 替换为你的图像路径
    if img is None:
        print("无法加载图像，请检查路径")
        return

    img_copy = img.copy()
    mask = np.zeros(img.shape[:2], np.uint8)

    # 创建窗口并绑定鼠标回调
    cv2.namedWindow('image')
    cv2.setMouseCallback('image', draw_rectangle)

    while True:
        cv2.imshow('image', img_copy)
        key = cv2.waitKey(1) & 0xFF

        # 按空格键开始GrabCut分割
        if key == 32 and rect != (0, 0, 0, 0):
            cv2.destroyWindow('image')
            grabcut_segmentation()
            # 绑定笔刷回调用于修正
            cv2.namedWindow('Segmentation Result')
            cv2.setMouseCallback('Segmentation Result', draw_brush)
            while True:
                cv2.imshow('Segmentation Result', img_copy)  # 显示修正图像
                key = cv2.waitKey(1) & 0xFF
                # 按'c'键重新GrabCut
                if key == ord('c'):
                    cv2.destroyWindow('Segmentation Result')
                    grabcut_segmentation()
                    cv2.namedWindow('Segmentation Result')
                    cv2.setMouseCallback('Segmentation Result', draw_brush)
                # 按其他键退出
                elif key != 255:  # 255是无按键
                    cv2.destroyAllWindows()
                    return
        # 按ESC退出
        elif key == 27:
            break

    cv2.destroyAllWindows()

if __name__ == '__main__':
    main()